#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
COCO 格式到单个 Labelme JSON 文件的转换工具
用于从 COCO 文件生成单个图片对应的 JSON 文件
"""

import json
from pathlib import Path
from typing import Dict, Optional, List


def coco_to_labelme_json(
    coco_file: Path,
    image_name: str,
    output_json: Path,
    image_width: int = 0,
    image_height: int = 0
) -> Optional[Dict]:
    """
    从 COCO 文件中提取指定图片的标注，生成 Labelme 格式的 JSON 文件
    
    Args:
        coco_file: COCO 格式的 JSON 文件路径
        image_name: 图片文件名
        output_json: 输出的 Labelme JSON 文件路径
        image_width: 图片宽度（如果为0，从COCO数据中获取）
        image_height: 图片高度（如果为0，从COCO数据中获取）
    
    Returns:
        Labelme 格式的数据字典，如果失败返回 None
    """
    try:
        # 加载 COCO 数据
        with open(coco_file, 'r', encoding='utf-8') as f:
            coco_data = json.load(f)
        
        # 查找图片信息（支持精确匹配和模糊匹配）
        # 规范化图片名称：去除路径，只保留文件名
        # COCO文件中的file_name可能是 "train_0045.jpg" 或 "annotations/train_0045.jpg" 或 "images/train_0045.jpg"
        # 统一处理为只比较文件名
        image_name_normalized = Path(image_name).name  # 确保只取文件名部分
        
        image_info = None
        for img in coco_data.get('images', []):
            coco_file_name = img.get('file_name', '')
            # 规范化COCO中的文件名（去除路径）
            coco_file_name_normalized = Path(coco_file_name).name
            
            # 精确匹配（比较规范化后的文件名）
            if coco_file_name_normalized == image_name_normalized:
                image_info = img
                break
            # 也尝试原始匹配（向后兼容）
            if coco_file_name == image_name or coco_file_name == image_name_normalized:
                image_info = img
                break
        
        if not image_info:
            # 输出调试信息
            print(f"⚠ 在COCO文件中未找到图片: {image_name}")
            # 列出前几个文件名供调试
            sample_files = [img.get('file_name', '') for img in coco_data.get('images', [])[:5]]
            if sample_files:
                print(f"   COCO文件中的示例文件名: {sample_files}")
            return None
        
        image_id = image_info['id']
        
        # 获取图片尺寸
        if image_width == 0:
            image_width = image_info.get('width', 0)
        if image_height == 0:
            image_height = image_info.get('height', 0)
        
        # 查找该图片的所有标注
        annotations = [
            ann for ann in coco_data.get('annotations', [])
            if ann.get('image_id') == image_id
        ]
        
        # 获取类别映射
        categories = {cat['id']: cat['name'] for cat in coco_data.get('categories', [])}
        
        # 转换为 Labelme 格式
        shapes = []
        for ann in annotations:
            bbox = ann.get('bbox', [])
            shape_type = 'rectangle'
            points = []
            
            if len(bbox) == 4:
                # 矩形框：Labelme 使用两个点（左上角和右下角）
                x, y, w, h = bbox
                points = [
                    [float(x), float(y)],           # 左上角
                    [float(x + w), float(y + h)]    # 右下角
                ]
            else:
                # 如果有 segmentation，使用 segmentation
                seg = ann.get('segmentation', [])
                if seg and len(seg) > 0:
                    shape_type = 'polygon'
                    if isinstance(seg[0], list):
                        seg_points = seg[0]
                        points = [[float(seg_points[i]), float(seg_points[i+1])] 
                                 for i in range(0, len(seg_points), 2)]
                    else:
                        points = [[float(seg[i]), float(seg[i+1])] 
                                 for i in range(0, len(seg), 2)]
                else:
                    continue
            
            if not points or len(points) < 2:
                continue
            
            category_id = ann.get('category_id')
            label = categories.get(category_id, f'category_{category_id}')
            
            shape = {
                'label': label,
                'points': points,
                'group_id': ann.get('id'),
                'shape_type': shape_type,
                'flags': {},
                'description': ''
            }
            shapes.append(shape)
        
        # 构建 Labelme 格式数据（兼容 Labelme 5.2.1）
        # imagePath 需要相对于 JSON 文件的位置，指向 images 目录
        # JSON在 annotations/xxx.json，图片在 images/xxx.jpg
        # 所以 imagePath 应该是 ../images/xxx.jpg（相对于annotations目录）
        image_path_relative = f"../images/{Path(image_name).name}"
        
        labelme_data = {
            'version': '5.2.1',  # 使用当前版本号
            'flags': {},
            'shapes': shapes,
            'imagePath': image_path_relative,  # 使用相对路径指向images目录
            # 不保存 imageData 字段（即使为 None 也不保存，减少文件大小）
            'imageWidth': image_width,
            'imageHeight': image_height
        }
        
        # 保存到文件
        output_json.parent.mkdir(parents=True, exist_ok=True)
        try:
            with open(output_json, 'w', encoding='utf-8') as f:
                json.dump(labelme_data, f, ensure_ascii=False, indent=2)
            print(f"[DEBUG] JSON文件已保存: {output_json.name}")
        except Exception as e:
            print(f"⚠ 保存JSON文件失败 {output_json}: {e}")
            import traceback
            traceback.print_exc()
            return None
        
        return labelme_data
        
    except Exception as e:
        print(f"⚠ 从 COCO 生成 Labelme JSON 失败: {e}")
        print(f"   图片: {image_name}")
        print(f"   输出: {output_json}")
        import traceback
        traceback.print_exc()
        return None


def get_coco_categories(coco_file: Path) -> List[str]:
    """
    从单个 COCO 文件中获取所有类别列表
    
    Args:
        coco_file: COCO 格式的 JSON 文件路径
    
    Returns:
        类别名称列表
    """
    try:
        with open(coco_file, 'r', encoding='utf-8') as f:
            coco_data = json.load(f)
        
        categories = []
        for cat in coco_data.get('categories', []):
            cat_name = cat.get('name', '')
            if cat_name:
                categories.append(cat_name)
        
        return sorted(categories)
        
    except Exception as e:
        print(f"⚠ 获取 COCO 类别失败: {e}")
        return []


def get_all_coco_categories(annotations_dir: Path) -> List[str]:
    """
    从所有 COCO 文件（instance_train.json 和 instance_val.json）中合并获取所有类别列表
    
    规范操作：从 instance_train.json 和 instance_val.json 中获取所有类别
    
    Args:
        annotations_dir: 标注文件目录
    
    Returns:
        合并后的类别名称列表（去重并排序）
    """
    all_categories = set()
    
    # 规范：从 instance_train.json 和 instance_val.json 中获取
    coco_files = [
        annotations_dir / 'instance_train.json',
        annotations_dir / 'instance_val.json',
    ]
    
    for coco_file in coco_files:
        if coco_file.exists():
            try:
                categories = get_coco_categories(coco_file)
                all_categories.update(categories)
                print(f"[DEBUG] 从 {coco_file.name} 加载了 {len(categories)} 个类别")
            except Exception as e:
                print(f"⚠ 从 {coco_file.name} 获取类别失败: {e}")
    
    # 返回排序后的类别列表
    return sorted(list(all_categories))

